import random
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity
from .models import MusicCollection  # 确保这个导入语句符合你的项目结构


class MusicRecommender:
    def __init__(self):
        self.user_music_matrix = None
        self.similarity_df = None
        self._prepare_data()

    def _prepare_data(self):
        """准备用户-音乐矩阵和用户相似度矩阵"""
        # 从MusicCollection模型加载数据
        data = list(MusicCollection.objects.all().values('user_id', 'music_id'))
        if not data:
            return []
        df = pd.DataFrame(data)

        # 构建用户-音乐矩阵
        self.user_music_matrix = df.pivot_table(index='user_id', columns='music_id', aggfunc='size', fill_value=0)

        # 计算用户间的相似度
        similarity_matrix = cosine_similarity(self.user_music_matrix)
        self.similarity_df = pd.DataFrame(similarity_matrix, index=self.user_music_matrix.index,
                                          columns=self.user_music_matrix.index)

    def recommend_music_for(self, user_id, top_n=3):
        """为指定用户推荐音乐"""
        if user_id not in self.similarity_df.index:
            return self.get_random_music(top_n)  # 如果用户ID不在矩阵中，返回随机推荐音乐

        # 找到最相似的用户
        similar_users = self.similarity_df[user_id].sort_values(ascending=False).index[1:top_n + 1]

        # 获取这些用户收藏的音乐
        similar_users_favorites = self.user_music_matrix.loc[similar_users].sum().sort_values(
            ascending=False).index.tolist()

        # 排除目标用户已收藏的音乐
        user_favorites = self.user_music_matrix.loc[user_id][self.user_music_matrix.loc[user_id] > 0].index.tolist()
        recommendations = [music_id for music_id in similar_users_favorites if music_id not in user_favorites][:top_n]

        # 如果没有推荐音乐，返回随机的音乐ID
        if not recommendations:
            recommendations = self.get_random_music(top_n)

        return recommendations

    def get_random_music(self, top_n):
        """获取随机的音乐ID"""
        all_music_ids = self.user_music_matrix.columns.tolist()
        print(all_music_ids)
        top_n = min(top_n, len(all_music_ids))
        print(top_n)
        return random.sample(all_music_ids, top_n)  # 随机选择 top_n 个音乐ID